# --------------------------------------------------------
# References:
# https://github.com/jxhe/unify-parameter-efficient-tuning
# --------------------------------------------------------
import torch
import torch.nn as nn

# --------------------------------------------------------
# References:
# timm: https://github.com/rwightman/pytorch-image-models/tree/master/timm
# DeiT: https://github.com/facebookresearch/deit
# MAE: https://github.com/facebookresearch/mae
# --------------------------------------------------------
from timm.models.vision_transformer import PatchEmbed
from timm.models.registry import register_model

from ..utils import initialize_vit_model
from ..vit import VisionTransformer as _VisionTransformer
from .._PEFTs import _PEFT_dict
from utils.toolkit import NamespaceDict


class VisionTransformer(_VisionTransformer):

    def __init__(
        self,
        img_size=224,
        patch_size=16,
        in_chans=3,
        num_classes=1000,
        embed_dim=768,
        depth=12,
        num_heads=12,
        mlp_ratio=4.0,
        qkv_bias=True,
        representation_size=None,
        distilled=False,
        drop_rate=0.0,
        attn_drop_rate=0.0,
        drop_path_rate=0.0,
        embed_layer=PatchEmbed,
        norm_layer=None,
        act_layer=None,
        weight_init="",
        global_pool=False,
        config: NamespaceDict = NamespaceDict(),
    ):
        super(VisionTransformer, self).__init__(
            img_size=img_size,
            patch_size=patch_size,
            in_chans=in_chans,
            num_classes=num_classes,
            embed_dim=embed_dim,
            depth=depth,
            num_heads=num_heads,
            mlp_ratio=mlp_ratio,
            qkv_bias=qkv_bias,
            representation_size=representation_size,
            distilled=distilled,
            drop_rate=drop_rate,
            attn_drop_rate=attn_drop_rate,
            drop_path_rate=drop_path_rate,
            embed_layer=embed_layer,
            norm_layer=norm_layer,
            act_layer=act_layer,
            weight_init=weight_init,
            global_pool=global_pool,
            config=config,
        )

        # ####### Adapter begins #########
        self.peft_func = _PEFT_dict[config.peft_name]
        self.cur_adapter = nn.ModuleList()
        for _ in range(len(self.blocks)):
            adapter = self.peft_func(
                self.config,
                embed_dim=self.config.embed_dim,
                bottleneck=self.config.ffn_rank,
                dropout=0.1,
                adapter_scalar=self.config.ffn_adapter_scalar,
                adapter_layernorm_option=self.config.ffn_adapter_layernorm_option,
            )
            self.cur_adapter.append(adapter)
        self.cur_adapter.requires_grad_(True)

    def forward_feats(self, x):
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)
        x = x + self.pos_embed
        x = self.pos_drop(x)

        for idx, blk in enumerate(self.blocks):
            x = blk(x, adapt=self.cur_adapter[idx])
        x = self.norm(x)
        return x

    def freeze(self):
        for param in self.parameters():
            param.requires_grad = False

        self.cur_adapter.requires_grad_(True)


@register_model
def vit_base_patch16_224_ssiat(pretrained=False, pretrained_cfg={}, **kwargs):
    del pretrained
    del pretrained_cfg
    return initialize_vit_model(
        "vit_base_patch16_224",
        VisionTransformer,
        **kwargs,
    )


@register_model
def vit_base_patch16_224_in21k_ssiat(
    pretrained=False, pretrained_cfg={}, **kwargs
):
    del pretrained
    del pretrained_cfg
    return initialize_vit_model(
        "vit_base_patch16_224_in21k",
        VisionTransformer,
        **kwargs,
    )
